import numpy as np
import matplotlib.pyplot as plt
from skimage import exposure as ex
import imageio
import sys
import cv2
import os


def he(img):
    if(len(img.shape)==2):      #gray
        outImg = ex.equalize_hist(img[:,:])*255 
    elif(len(img.shape)==3):    #RGB
        outImg = np.zeros((img.shape[0],img.shape[1],3))
        for channel in range(img.shape[2]):
            outImg[:, :, channel] = ex.equalize_hist(img[:, :, channel])*255

    outImg[outImg>255] = 255
    outImg[outImg<0] = 0
    return outImg.astype(np.uint8)


def main():
    path = './testdata/'
    images = [x for x in os.listdir(path)]
    for img_name in images:
        img_path = path + img_name
        img = cv2.imread(img_path)
        # img = imageio.imread(img_path)
        name = img_name.split('.')[0]
        result = he(img)
        cv2.imwrite(name + '-ressult' + '.bmp', result)


if __name__ == '__main__':
    main()